{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from sklearn.tree import DecisionTreeClassifier  # Import Decision Tree Classifier\n",
    "from sklearn.model_selection import train_test_split  # Import train_test_split function\n",
    "from sklearn import metrics  # Import scikit-learn metrics module for accuracy calculation\n",
    "from sklearn import tree\n",
    "from scipy.io import arff\n",
    "import re\n",
    "import sys\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def isLeaf(src_str):\n",
    "    r = re.match(\"[0-9]+ \\[label=\\\"\", src_str)\n",
    "    s = re.search(\"\\\\\\\\n\", src_str)\n",
    "    if (r != None and s == None):\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "def printClassName(src_str, dt_file):\n",
    "    r = re.search(\"[0-9]+ \\[label=\\\"\", src_str)\n",
    "    s = re.search(\"\\\"\\] ;\", src_str)\n",
    "    t = re.search(\"\\\\\\\\n\", src_str)\n",
    "    if (r != None and s != None and t == None):\n",
    "        print(src_str[r.end():s.start()], end=\"\", file=dt_file)\n",
    "\n",
    "\n",
    "def printCondition(src_str, dt_file):\n",
    "    r = re.search(\"[0-9]+ \\[label=\\\"\", src_str)\n",
    "    s = re.search(\"\\\\\\\\n\", src_str)\n",
    "    if (r != None and s != None):\n",
    "        print(src_str[r.end():s.start()], end=\"\", file=dt_file)\n",
    "\n",
    "\n",
    "def getNodeNum(src_str):\n",
    "    r = re.match(\"[0-9]+\", src_str)\n",
    "    if (r != None):\n",
    "        return int(r.group(0))\n",
    "    else:\n",
    "        return \"\"\n",
    "\n",
    "\n",
    "def getNextLineIndex(src_list, node_num):\n",
    "    tmp = []\n",
    "    i = len(src_list) - 1\n",
    "    for line in reversed(src_list):\n",
    "        if (getNodeNum(line) == int(node_num)):\n",
    "            return i\n",
    "        i -= 1\n",
    "    return -1\n",
    "\n",
    "\n",
    "def isNodeInfo(src_str):\n",
    "    if (re.match('[0-9]+ -> [0-9]+[ ]+;', src_str) != None):\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "def oppOperator(src_str):\n",
    "    src_str = src_str.replace(\"<=\", \">\")\n",
    "    src_str = src_str.replace(\">=\", \"<\")\n",
    "    return src_str\n",
    "\n",
    "\n",
    "def formatTree(line, indent, dt_file):\n",
    "    # Init to zeros, will store n_nodes and n_leaves\n",
    "    size = np.array([0, 0])\n",
    "    # If the first line is just node connection info line, skip it\n",
    "    if (isNodeInfo(line[0])):\n",
    "        line = line[1:]\n",
    "\n",
    "    # Add this node\n",
    "    size[0] += 1\n",
    "\n",
    "    # If the first line is a leaf, print its name with \\n, otherwise only \\n\n",
    "    if (isLeaf(line[0])):\n",
    "        print(\": \", end=\"\", file=dt_file)\n",
    "        printClassName(line[0], dt_file)\n",
    "        print(\"\", file=dt_file)\n",
    "        return np.array([1, 1])\n",
    "    else:\n",
    "        print(\"\", file=dt_file)\n",
    "\n",
    "    nIndex = getNodeNum(line[0])  # Get node index\n",
    "    splitIndex = getNextLineIndex(line[1:], nIndex)  # Get split index\n",
    "\n",
    "    if (len(line[1:splitIndex]) > 0):\n",
    "        # Print original condition\n",
    "        print(\"|   \" * indent, end=\"\", file=dt_file)\n",
    "        printCondition(line[0], dt_file)\n",
    "        size += formatTree(line[1:splitIndex], indent + 1, dt_file)  # Call recursively for the first part of original tree\n",
    "\n",
    "    if (len(line[splitIndex - 1:]) > 0):\n",
    "        # Print opposite condition\n",
    "        print(\"|   \" * indent, end=\"\", file=dt_file)\n",
    "        printCondition(oppOperator(line[0]), dt_file)\n",
    "        size += formatTree(line[splitIndex - 1:], indent + 1, dt_file)  # Call recursively for the second part of original tree\n",
    "\n",
    "    return size\n",
    "\n",
    "\n",
    "def printTree(dot_tree, dt_file):\n",
    "    new_tree = []\n",
    "\n",
    "    # Preprocess the tree\n",
    "    for line in dot_tree.split(\"\\n\"):\n",
    "        r = re.search(\"[0-9]+\\\\\\\\n\\[([0-9]+[,]?[ ]?)+\\]\\\\\\\\n\", line)\n",
    "        s = re.search(\"\\[labeldistance=[0-9]+\\.?[0-9]*, labelangle=-?[0-9]+, headlabel=\\\"(False|True)\\\"\\]\", line)\n",
    "        if (r != None):\n",
    "            line = line[:r.start()] + line[r.end():]\n",
    "        if (s != None):\n",
    "            line = line[:s.start()] + line[s.end():]\n",
    "        new_tree.append(line)\n",
    "\n",
    "    # Print in Weka format\n",
    "    n_nodes, n_leaves = formatTree(new_tree[2:-1], 0, dt_file)\n",
    "\n",
    "    print('\\nNumber of Leaves  : \\t', n_leaves, file=dt_file)\n",
    "    print('\\nSize of the Tree : \\t', n_nodes, file=dt_file)\n",
    "\n",
    "    return n_nodes, n_leaves\n",
    "\n",
    "\n",
    "def generateDecisionTree(arff_filename, dectree_filename):\n",
    "    data = arff.loadarff(arff_filename)  # <- Write desired file here\n",
    "    data_set = pd.DataFrame(data[0])\n",
    "    data_set['class'] = data_set['class'].str.decode('ASCII')\n",
    "    col_names = list(data_set)\n",
    "\n",
    "    feature_cols = col_names[:-1]\n",
    "    class_name = list(set(data_set[col_names[-1]]))\n",
    "\n",
    "    X = data_set[feature_cols]  # Features\n",
    "    y = data_set[col_names[-1]]  # Target variable\n",
    "\n",
    "    # Split dataset into training set and test set\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3,\n",
    "                                                        random_state=1)  # 70% training and 30% test\n",
    "\n",
    "    # Create Decision Tree classifer object\n",
    "    clf = tree.DecisionTreeClassifier(random_state=1)  # You can specify the max depth by passing argument for example: max_depth=3\n",
    "\n",
    "    # Train Decision Tree Classifer\n",
    "    clf = clf.fit(X_train, y_train)\n",
    "\n",
    "    # Predict the response for test dataset\n",
    "    y_pred = clf.predict(X_test)\n",
    "\n",
    "    # Model Accuracy, how often is the classifier correct?\n",
    "    dectree_accuracy = metrics.accuracy_score(y_test, y_pred)\n",
    "    print(\"Accuracy:\", dectree_accuracy)\n",
    "\n",
    "    dot_tree = tree.export_graphviz(clf, out_file=None, class_names=clf.classes_, label=\"none\", impurity=False,\n",
    "                                    feature_names=feature_cols)\n",
    "\n",
    "    dt_file = open(dectree_filename, \"w\")\n",
    "    n_nodes, n_leaves = printTree(dot_tree, dt_file)\n",
    "    dt_file.close()\n",
    "\n",
    "    printTree(dot_tree, sys.stdout)\n",
    "\n",
    "    return dectree_accuracy, n_nodes\n",
    "\n",
    "\n",
    "def generate_subset_of_ARFF(arff_filename, arff_subset_filename, classes_subset):\n",
    "    with open(arff_filename, \"r\") as input:\n",
    "        with open(arff_subset_filename, \"w\") as output:\n",
    "            data_section_started = False\n",
    "            for line in input:\n",
    "                if line.startswith('@attribute class {'):\n",
    "                    new_line_class = '@attribute class {'\n",
    "                    for i_class in range(len(classes_subset)):\n",
    "                        if i_class > 0:\n",
    "                            new_line_class += ', '\n",
    "                        new_line_class += classes_subset[i_class]\n",
    "                    new_line_class += '}\\r\\n'\n",
    "                    output.write(new_line_class)\n",
    "                else:\n",
    "                    if line.startswith('@data'):\n",
    "                        data_section_started = True\n",
    "                        output.write(line)\n",
    "\n",
    "                    if (data_section_started == False):\n",
    "                        output.write(line)\n",
    "                    else:\n",
    "                        class_in_line = line.split(', ')[-1]\n",
    "                        if (class_in_line.rstrip() in classes_subset):\n",
    "                            output.write(line)\n",
    "    input.close()\n",
    "    output.close()\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "arff_filename = \"test.arff\"\n",
    "dectree_filename = \"dectree.txt\"\n",
    "dectree_accuracy, dectree_nodes = generateDecisionTree(arff_filename, dectree_filename)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
